// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/base/indexed_priority_queue.h>

#include <melon/benchmark.h>
#include <melon/random.h>
#include <melon/init/init.h>

#define POLLUX_BENCHMARK(_type, _name, ...) \
  [[maybe_unused]] _type _name(MELON_PP_STRINGIZE(_name), __VA_ARGS__)

namespace kumo::pollux {
namespace {

template <typename T>
class IndexedPriorityQueueBenchmark {
 public:
  IndexedPriorityQueueBenchmark(const char* name, int numValues)
      : values_(numValues), priorities_(2 * numValues) {
    generateValues();
    for (int i = 0; i < 2 * numValues; ++i) {
      priorities_[i] = melon::Random::rand32();
    }
    melon::addBenchmark(__FILE__, fmt::format("{}_add", name), [this] {
      IndexedPriorityQueue<T, true> queue;
      return add(queue);
    });
    melon::addBenchmark(__FILE__, fmt::format("{}_update", name), [this] {
      IndexedPriorityQueue<T, true> queue;
      BENCHMARK_SUSPEND {
        add(queue);
      }
      return update(queue);
    });
    melon::addBenchmark(__FILE__, fmt::format("{}_pop", name), [this] {
      IndexedPriorityQueue<T, true> queue;
      BENCHMARK_SUSPEND {
        add(queue);
        update(queue);
      }
      return pop(queue);
    });
  }

 private:
  void generateValues();

  unsigned add(IndexedPriorityQueue<T, true>& queue) const {
    for (int i = 0; i < values_.size(); ++i) {
      queue.addOrUpdate(values_[i], priorities_[i]);
    }
    return values_.size();
  }

  unsigned update(IndexedPriorityQueue<T, true>& queue) const {
    for (int i = 0; i < values_.size(); ++i) {
      queue.addOrUpdate(values_[i], priorities_[i + values_.size()]);
    }
    return values_.size();
  }

  unsigned pop(IndexedPriorityQueue<T, true>& queue) const {
    while (!queue.empty()) {
      queue.pop();
    }
    return values_.size();
  }

  std::vector<T> values_;
  std::vector<uint32_t> priorities_;
};

template <>
void IndexedPriorityQueueBenchmark<int64_t>::generateValues() {
  std::iota(values_.begin(), values_.end(), 0);
}

} // namespace
} // namespace kumo::pollux

int main(int argc, char* argv[]) {
  using namespace kumo::pollux;
  melon::Init melonInit(&argc, &argv);
  POLLUX_BENCHMARK(IndexedPriorityQueueBenchmark<int64_t>, int64_1000, 1000);
  POLLUX_BENCHMARK(IndexedPriorityQueueBenchmark<int64_t>, int64_10000, 10'000);
  POLLUX_BENCHMARK(
      IndexedPriorityQueueBenchmark<int64_t>, int64_100000, 100'000);
  POLLUX_BENCHMARK(
      IndexedPriorityQueueBenchmark<int64_t>, int64_1000000, 1'000'000);
  melon::runBenchmarks();
  return 0;
}
